AWS Step Functions から直接 AWS Batch のジョブ実行開始を試してみた

Step Functionsから直接AWS Batchを実行できる!

#AWS Step Functions

#AWS Batch

#HPC

#AWS Lambda

#AWS

大村保貴

2023.07.17

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

AWS Batch でジョブを実行する前処理としてLambda の実行結果（出力した値）を使って AWS Batch でジョブを実行したいです。

このような制御する場合は何かしらのワークフローが必要になるため、Step Functions を利用するのが一般的かと思います。

AWS Batch のジョブ実行をキックする Lambda を書くつもりでいたのですが、Workflow Studio で Batch と検索するとSubmitJob という「顧客（私）が本当に必要だったもの」と思われるものが用意されていました。

SubmitJobの使用感確認のために Step Fuctions のワークフローでSubmitJobを利用して、AWS Batch のジョブを実行できるのかをやってみました。

各種設定

Lambda の実行結果（出力した値）を AWS Batch のジョブ実行時に渡す検証として以下のワークフローを作成します。Lambda と AWS Batch、そして、Step Functions で必要になった主要な設定を記載します。

Lambda の設定

Lambda 関数はreturnに固定値を返すだけの関数を用意しました。

def lambda_handler(event, context):
    return {
        'statusCode': 200,
        'target': '/path/sample.txt'
    }

テスト実行した結果は以下になります。targetキーの値を AWS Batch に引き渡します。

{
  "statusCode": 200,
  "target": "/path/sample.txt"
}

AWW Batch の設定

パラメーターMessageを用意し、Step Functions からMessageの値を書き換えます。ジョブを実行するとecho Ref::MessageでMessageの内容を標準出力に返します。デフォルト値としてhiを入力しています。

AWS BatchでRef::はパラメーターの値を使うプレースホルダーの記述です。

ジョブ定義のパラメータ - AWS Batch

ジョブ定義全文

{
    "jobDefinitions": [
        {
            "jobDefinitionName": "fagate-sfn",
            "jobDefinitionArn": "arn:aws:batch:ap-northeast-1:123456789012:job-definition/fagate-sfn:1",
            "revision": 1,
            "status": "ACTIVE",
            "type": "container",
            "parameters": {
                "Message": "hi"
            },
            "containerProperties": {
                "image": "public.ecr.aws/amazonlinux/amazonlinux:latest",
                "command": [
                    "echo",
                    "Ref::Message"
                ],
                "jobRoleArn": "arn:aws:iam::123456789012:role/batch-test-BatchTaskExecutionRole",
                "executionRoleArn": "arn:aws:iam::123456789012:role/batch-test-BatchTaskExecutionRole",
                "volumes": [],
                "environment": [],
                "mountPoints": [],
                "ulimits": [],
                "resourceRequirements": [
                    {
                        "value": "1.0",
                        "type": "VCPU"
                    },
                    {
                        "value": "2048",
                        "type": "MEMORY"
                    }
                ],
                "logConfiguration": {
                    "logDriver": "awslogs",
                    "options": {},
                    "secretOptions": []
                },
                "secrets": [],
                "networkConfiguration": {
                    "assignPublicIp": "ENABLED"
                },
                "fargatePlatformConfiguration": {
                    "platformVersion": "LATEST"
                }
            },
            "tags": {},
            "platformCapabilities": [
                "FARGATE"
            ],
            "containerOrchestrationType": "ECS"
        }
    ]
}

ジョブ定義の JSON 出力方法

aws batch describe-job-definitionsコマンドで出力できます。

`aws batch describe-job-definitions --job-definition-name [INPUT_NAME]`

Step Functions の設定

ワークフローは Lambda を実行した後に、AWS Batch でジョブを実行する直列なフローを作成しました。

今回の検証ポイント

Lambda の実行結果を$とし、SubmitJob 実行のパラメーターで$.targetとしています。Lambda の実行結果targetの値がMessageパラメーターの値になっています。

AWS Batch のジョブ定義ではMessageの値を echo コマンドで標準出力に返すだけですので、Lambda の実行結果であるtargetの値が AWS Batch で実行したコンテナの標準出力に表示される寸法です。

InputPath、パラメータ、ResultSelector - AWS Step Functions

ここが今回の検証したかった SubmitJob へ必要なパラメーターの渡し方と、AWS Batch のジョブ定義で Step Functions のワークフローから受け取った値を利用する方法です。

ワークフロー定義全文

{
  "Comment": "A description of my state machine",
  "StartAt": "Pass the result",
  "States": {
    "Pass the result": {
      "Type": "Task",
      "Resource": "arn:aws:states:::lambda:invoke",
      "OutputPath": "$.Payload",
      "Parameters": {
        "Payload.$": "$",
        "FunctionName": "arn:aws:lambda:ap-northeast-1:123456789012:function:test-sfn-in-out-value:$LATEST"
      },
      "Retry": [
        {
          "ErrorEquals": [
            "Lambda.ServiceException",
            "Lambda.AWSLambdaException",
            "Lambda.SdkClientException",
            "Lambda.TooManyRequestsException"
          ],
          "IntervalSeconds": 2,
          "MaxAttempts": 6,
          "BackoffRate": 2
        }
      ],
      "Next": "Batch SubmitJob"
    },
    "Batch SubmitJob": {
      "Type": "Task",
      "Resource": "arn:aws:states:::batch:submitJob.sync",
      "Parameters": {
        "JobName": "fargate-test",
        "JobDefinition": "arn:aws:batch:ap-northeast-1:123456789012:job-definition/fagate-sfn:1",
        "JobQueue": "arn:aws:batch:ap-northeast-1:123456789012:job-queue/fargate-queue",
        "Parameters": {
          "Message.$": "$.target"
        }
      },
      "End": true
    }
  }
}

以上が主要な設定内容でした。

ワークフローを実行してみる

ワークフローを実行し正常に終了しました。結果を確認します。

Lambda の実行結果

Lambda の出力タブにはtargetの値を固定値（/path/sample.txt）で返した内容を確認できます。

AWS Batch の実行結果

AWS Batch の入力タブには Lambda の出力内容がそのまま渡されています。

出力タブの内容は画像で収まらないので出力結果（JSON形式）を書き出します。

ジョブ定義のパラメーターでプレースホルダーとしていた$.targetの箇所が/path/sample.txtに置き換わっており期待した結果になっています。

{
  "Attempts": [
    {
      "Container": {
        "ExitCode": 0,
        "LogStreamName": "fagate-sfn/default/a07ea7aef82e4a04ba5237ce4f320c09",
        "NetworkInterfaces": [
          {
            "AttachmentId": "679f1438-636b-4fd7-a540-410057a2779f",
            "PrivateIpv4Address": "10.0.1.87"
          }
        ],
        "TaskArn": "arn:aws:ecs:ap-northeast-1:123456789012:task/AWSBatch-fargate-1a-846f808b-7caf-35a9-9192-292884f9bd46/a07ea7aef82e4a04ba5237ce4f320c09"
      },
      "StartedAt": 1689557537640,
      "StatusReason": "Essential container in task exited",
      "StoppedAt": 1689557560856
    }
  ],
  "Container": {
    "Command": [
      "echo",
      "/path/sample.txt"
    ],
    "Environment": [
      {
        "Name": "MANAGED_BY_AWS",
        "Value": "STARTED_BY_STEP_FUNCTIONS"
      }
    ],
    "ExecutionRoleArn": "arn:aws:iam::123456789012:role/batch-test-BatchTaskExecutionRole",
    "ExitCode": 0,
    "FargatePlatformConfiguration": {
      "PlatformVersion": "LATEST"
    },
    "Image": "public.ecr.aws/amazonlinux/amazonlinux:latest",
    "JobRoleArn": "arn:aws:iam::123456789012:role/batch-test-BatchTaskExecutionRole",
    "LogConfiguration": {
      "LogDriver": "awslogs",
      "Options": {},
      "SecretOptions": []
    },
    "LogStreamName": "fagate-sfn/default/a07ea7aef82e4a04ba5237ce4f320c09",
    "MountPoints": [],
    "NetworkConfiguration": {
      "AssignPublicIp": "ENABLED"
    },
    "NetworkInterfaces": [
      {
        "AttachmentId": "679f1438-636b-4fd7-a540-410057a2779f",
        "PrivateIpv4Address": "10.0.1.87"
      }
    ],
    "ResourceRequirements": [
      {
        "Type": "VCPU",
        "Value": "1.0"
      },
      {
        "Type": "MEMORY",
        "Value": "2048"
      }
    ],
    "Secrets": [],
    "TaskArn": "arn:aws:ecs:ap-northeast-1:123456789012:task/AWSBatch-fargate-1a-846f808b-7caf-35a9-9192-292884f9bd46/a07ea7aef82e4a04ba5237ce4f320c09",
    "Ulimits": [],
    "Volumes": []
  },
  "CreatedAt": 1689557509347,
  "DependsOn": [],
  "EksAttempts": [],
  "JobArn": "arn:aws:batch:ap-northeast-1:123456789012:job/cb13b2a9-0bd2-490f-b3cb-a24948849bde",
  "JobDefinition": "arn:aws:batch:ap-northeast-1:123456789012:job-definition/fagate-sfn:1",
  "JobId": "cb13b2a9-0bd2-490f-b3cb-a24948849bde",
  "JobName": "fargate-test",
  "JobQueue": "arn:aws:batch:ap-northeast-1:123456789012:job-queue/fargate-queue",
  "Parameters": {
    "Message": "/path/sample.txt"
  },
  "PlatformCapabilities": [
    "FARGATE"
  ],
  "StartedAt": 1689557537640,
  "Status": "SUCCEEDED",
  "StatusReason": "Essential container in task exited",
  "StoppedAt": 1689557560856,
  "Tags": {
    "resourceArn": "arn:aws:batch:ap-northeast-1:123456789012:job/cb13b2a9-0bd2-490f-b3cb-a24948849bde"
  }
}

AWS Batch で実行したジョブの内容は echo コマンドで、パラメーターMessageの値（ここで期待するのは/path/sample.txt）を標準出力に返すといった内容でした。

実際の AWS Batch の実行ログを確認します。Step Functions の実行結果から AWS Batch のログを確認するリンクが貼られています。

AWS Batch で実行したジョブのログを確認すると/path/sample.txtの文字列を確認できました。

これで Lambda の実行結果の値を Step Functions の SumitJob を利用して AWS Batch のジョブ実行時に利用する動作確認が取れました。

おわりに

AWS Batch のジョブを実行開始するだけのために Lambda を用意しなくても良い時代になっていました。込み入ったジョブの実行の仕方だと引き続き Lambda から AWS Batch の SDK を使って API を叩くことになりますが、まずは Step Fuctions の SubmitJob を利用して実行できないか検討してみると良いのでないでしょうか。Lambda を省略できるならより運用負荷を削減できますよ。

他にも Lambda なしで実行できるものないか探してみてください。